/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "OpenSparseBLAS_func.h"
#include <iostream>
#include <cstddef>
#include <complex>
#include <type_traits>
#include <cassert>

// -----------------------------------------------------------------------------
// 通用 CSC 稀疏矩阵 × 稠密矩阵：
//   C = beta * C + alpha * op(A) * B
// A  : mA × kA，CSC  格式 (columnPtr, rowIdx, valA)
// B  : kA × n ，列主序，步长 ldb
// C  : mC × n ，列主序，步长 ldc
//
// 其中：
//   ─ NONTRANS：op(A) =  A          (m = mA, k = kA)
//   ─ TRANS   ：op(A) =  Aᵀ        (m = kA, k = mA)
//   ─ CONJTRANS（复数专用）：op(A) = Aᴴ
//
// 只修正正确性，不做性能优化。
// -----------------------------------------------------------------------------
template <typename T>
void OpenSpB_cscmm(
    OpenSpB_trans_type transa,
    OpenSpB_base_type  index_base,
    int  m,                   // op(A) 的行数
    int  n,                   // B/C 的列数
    int  k,                   // op(A) 的列数
    const T& alpha,
    const int* colPtr,        // kA+1
    const int* rowIdx,        // nnz
    const T*  valA,           // nnz
    const T*  B, int ldb,
    const T&  beta,
    T*  C, int ldc
)
{
    // -------------------------------------------------------------------------
    // 基础合法性检查（仅覆盖越界相关的关键项，其他留给上层保证）
    // -------------------------------------------------------------------------
    if (m <= 0 || n <= 0 || k <= 0 || ldb <= 0 || ldc <= 0 ||
        !colPtr || !rowIdx || !valA || !B || !C)
    {
        std::cerr << "error: illegal input (dimensions or null pointer)\n";
        return;
    }

    if (!(index_base == OPENSPBLAS_ZERO_BASE ||
          index_base == OPENSPBLAS_ONE_BASE))
    {
        std::cerr << "error: unsupported index_base (must be 0- or 1-based)\n";
        return;
    }

    const int base = (index_base == OPENSPBLAS_ONE_BASE) ? 1 : 0;

    // colPtr[0] 必须等于 base
    assert(colPtr[0] == base &&
           "colPtr[0] must be 0 (zero-based) or 1 (one-based)");

    // -------------------------------------------------------------------------
    // 先处理 C ← beta * C
    // -------------------------------------------------------------------------
    if (beta == T{}) {                        // beta == 0 → 直接清零
        for (int j = 0; j < n; ++j)
            for (int i = 0; i < m; ++i)
                C[i + j * ldc] = T{};
    }
    else if (beta != T{1}) {                  // beta ≠ 1 → 缩放
        for (int j = 0; j < n; ++j)
            for (int i = 0; i < m; ++i)
                C[i + j * ldc] *= beta;
    }

    if (alpha == T{}) return;                 // alpha == 0，计算可结束

    // -------------------------------------------------------------------------
    // 不转置：C(m×n) += alpha * A(m×k) * B(k×n)
    // -------------------------------------------------------------------------
    if (transa == OPENSPBLAS_NO_TRANS)
    {
        for (int colA = 0; colA < k; ++colA)          // 遍历 A 的列 (0 .. kA-1)
        {
            const int col_start = colPtr[colA]   - base;
            const int col_end   = colPtr[colA+1] - base;

            // 对应 B 的第 colA 行（列主序，行距 = 1）
            const T* Bj = B + colA;

            for (int p = col_start; p < col_end; ++p) // 遍历该列的非零
            {
                const int rowA = rowIdx[p] - base;    // A 的行号
                const T   aij  = alpha * valA[p];

                T* Ci = C + rowA;                     // C 的 (rowA, :)
                for (int colB = 0; colB < n; ++colB)
                    Ci[colB * ldc] += aij * Bj[colB * ldb];
            }
        }
    }
    // -------------------------------------------------------------------------
    // 转置 / 共轭转置：C(m×n) += alpha * Aᵀ(m×k) * B(k×n)
    //
    // 此时：
    //   m = kA (A 的列数)，k = mA (A 的行数)
    //   colPtr 长度 = kA + 1 = m + 1
    //   外层应遍历 **A 的列** → 索引范围 0 .. kA-1 = m-1  (FIX)
    // -------------------------------------------------------------------------
    else if (transa == OPENSPBLAS_TRANS || transa == OPENSPBLAS_CONJ_TRANS)
    {
        const bool conj = (transa == OPENSPBLAS_CONJ_TRANS);

        for (int colA = 0; colA < m; ++colA)          // FIX: 用 m 而非 k
        {
            const int col_start = colPtr[colA]   - base;
            const int col_end   = colPtr[colA+1] - base;

            for (int p = col_start; p < col_end; ++p)
            {
                const int rowA_raw = rowIdx[p] - base; // 0 .. mA-1 (= k-1)
                assert(rowA_raw >= 0 && rowA_raw < k);

                // 复数共轭
                T aij = valA[p];
                if (conj) {
                    if constexpr (std::is_same_v<T, std::complex<float>> ||
                                  std::is_same_v<T, std::complex<double>>) {
                        aij = std::conj(aij);
                    }
                }
                aij *= alpha;

                // B 的 rowA_raw 行        (行距 = 1)
                const T* Bj = B + rowA_raw;
                // C 的 colA    行 (因为 C 是列主序，同样行距 = 1)
                T* Ci = C + colA;

                for (int colB = 0; colB < n; ++colB)
                    Ci[colB * ldc] += aij * Bj[colB * ldb];
            }
        }
    }
    else
    {
        std::cerr << "error: unsupported transa type\n";
    }
}

// -----------------------------------------------------------------------------
// s / d / c / z 外壳函数（仅转调模版核心函数）
// -----------------------------------------------------------------------------
void OpenSpB_scscmm(
    OpenSpB_trans_type transa,
    OpenSpB_base_type index_base,
    const int* m, const int* n, const int* k,
    const float* alpha,
    const int* columnPtra, const int* rowIndexa, const float* vala,
    const float* b, int ldb,
    const float* beta,
    float* c, int ldc)
{
    OpenSpB_cscmm<float>(
        transa, index_base,
        *m, *n, *k,
        *alpha,
        columnPtra, rowIndexa, vala,
        b, ldb,
        *beta,
        c, ldc
    );
}

void OpenSpB_dcscmm(
    OpenSpB_trans_type transa,
    OpenSpB_base_type index_base,
    const int* m, const int* n, const int* k,
    const double* alpha,
    const int* columnPtra, const int* rowIndexa, const double* vala,
    const double* b, int ldb,
    const double* beta,
    double* c, int ldc)
{
    OpenSpB_cscmm<double>(
        transa, index_base,
        *m, *n, *k,
        *alpha,
        columnPtra, rowIndexa, vala,
        b, ldb,
        *beta,
        c, ldc
    );
}

void OpenSpB_ccscmm(
    OpenSpB_trans_type transa,
    OpenSpB_base_type index_base,
    const int* m, const int* n, const int* k,
    const void* alpha,
    const int* columnPtra, const int* rowIndexa, const void* vala,
    const void* b, int ldb,
    const void* beta,
    void* c, int ldc)
{
    const auto a   = *static_cast<const std::complex<float>*>(alpha);
    const auto bta = *static_cast<const std::complex<float>*>(beta);

    OpenSpB_cscmm<std::complex<float>>(
        transa, index_base,
        *m, *n, *k,
        a,
        columnPtra, rowIndexa,
        static_cast<const std::complex<float>*>(vala),
        static_cast<const std::complex<float>*>(b), ldb,
        bta,
        static_cast<std::complex<float>*>(c), ldc
    );
}

void OpenSpB_zcscmm(
    OpenSpB_trans_type transa,
    OpenSpB_base_type index_base,
    const int* m, const int* n, const int* k,
    const void* alpha,
    const int* columnPtra, const int* rowIndexa, const void* vala,
    const void* b, int ldb,
    const void* beta,
    void* c, int ldc)
{
    const auto a   = *static_cast<const std::complex<double>*>(alpha);
    const auto bta = *static_cast<const std::complex<double>*>(beta);

    OpenSpB_cscmm<std::complex<double>>(
        transa, index_base,
        *m, *n, *k,
        a,
        columnPtra, rowIndexa,
        static_cast<const std::complex<double>*>(vala),
        static_cast<const std::complex<double>*>(b), ldb,
        bta,
        static_cast<std::complex<double>*>(c), ldc
    );
}
